package com.mayee.me.gateway.filter;

import io.netty.buffer.ByteBufAllocator;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.net.URLEncoder;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;

/**
 * @program: micro-service
 * @description: 缓存请求参数，只支持缓存下面 2 种 MediaType
 * APPLICATION_JSON--Json数据
 * APPLICATION_FORM_URLENCODED--FormData表单数据
 *
 * 缓存 body 时，不能在 filter 内部直接缓存，需要按响应式处理，在异步操作路途上缓存 body ，由于 body 只能读取一次，所以需要重新封装 request
 * 和 exchange 保证请求正常传递到下游;
 * 缓存 formData 时，formData 也只能读取一次，也需要重新封装 request 和 exchange。注意，如果对 formData 进行修改，必须重新定义 header
 * 中的 content-length，以保证传输数据大小一致。
 *
 * @author: Bobby.Ma
 * @create: 2021-01-09 10:17
 **/
@Slf4j
@Component
public class GatewayContextFilter implements GlobalFilter, Ordered {
    /*
     * default HttpMessageReader
     */
    private static final List<HttpMessageReader<?>> MESSAGE_READERS = HandlerStrategies.withDefaults().messageReaders();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        GatewayContext gatewayContext = new GatewayContext();
        URI uri = request.getURI();
        gatewayContext.setPath(uri.getPath());
        log.info("query---> :{}",uri.getQuery());
        exchange.getAttributes().put(GatewayContext.GATEWAY_CONTEXT,gatewayContext);
        HttpHeaders headers = request.getHeaders();
        MediaType contentType = headers.getContentType();
        log.info("start-------------------------------------------------");
        log.info("HttpMethod:{},Url:{}", request.getMethod(), uri.getPath());
        if (Objects.equals(request.getMethod(), HttpMethod.GET)){
            log.info("end-------------------------------------------------");
        }
        if (Objects.equals(request.getMethod(), HttpMethod.POST)){
            Mono<Void> Mono = null;
            if (Objects.equals(contentType,MediaType.APPLICATION_JSON)){
                Mono = readBodyData(exchange, chain, gatewayContext);
            }
            if (Objects.equals(contentType,MediaType.APPLICATION_FORM_URLENCODED)){
                Mono = readFormData(exchange, chain, gatewayContext);
            }
            return Mono;
        }
        return chain.filter(exchange);
    }

    private Mono<Void> readBodyData(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        final ServerHttpRequest request = exchange.getRequest();
        /*
         * join the body
         */
        return DataBufferUtils.join(exchange.getRequest().getBody()).flatMap(dataBuffer -> {
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer);
                    Flux<DataBuffer> bufferFlux = Flux.defer(() -> {
                        DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                        DataBufferUtils.retain(buffer);
                        return Mono.just(buffer);
                    });
                    RequestDecorator decorator = new RequestDecorator(request,bufferFlux);
                    /*
                     * mutate exchange with new ServerHttpRequest
                     */
                    ServerWebExchange mutatedExchange = exchange.mutate().request(decorator).build();
                    /*
                     * read body string with default messageReaders
                     */
                    return ServerRequest.create(mutatedExchange, MESSAGE_READERS)
                            .bodyToMono(String.class)
                            .doOnNext(objectValue -> {
                                log.info("PostBody:{}", objectValue);
                                log.info("end-------------------------------------------------");
                                gatewayContext.setBody(objectValue);
                            }).then(chain.filter(mutatedExchange));
                });
    }

    private Mono<Void> readFormData(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        final ServerHttpRequest request = exchange.getRequest();
        return exchange.getFormData().doOnNext(multiValueMap -> {
                    gatewayContext.setFormData(multiValueMap);
                    log.info("Post x-www-form-urlencoded:{}", multiValueMap);
                    log.info("end-------------------------------------------------");
                }).then(Mono.defer(() -> {
                    MultiValueMap<String, String> formData = gatewayContext.getFormData();
                    /*
                     * formData is empty just return
                     */
                    if (CollectionUtils.isEmpty(formData)){
                        return chain.filter(exchange);
                    }
                    Charset charset = Optional.ofNullable(request.getHeaders().getContentType())
                            .map(MediaType::getCharset)
                            .orElse(StandardCharsets.UTF_8);
                    /*
                     * repackage form data
                     */
                    String packageFormData = packageFormData(formData, charset.name());
                    /*
                     * get data bytes
                     */
                    byte[] bodyBytes = packageFormData.getBytes(charset);
                    Flux<DataBuffer> bufferFlux = DataBufferUtils.read(
                            new ByteArrayResource(bodyBytes),
                            new NettyDataBufferFactory(ByteBufAllocator.DEFAULT),
                            bodyBytes.length);
                    RequestDecorator decorator = new RequestDecorator(request,bufferFlux);
                    ServerWebExchange mutateExchange = exchange.mutate().request(decorator).build();
                    return chain.filter(mutateExchange);
                }));
    }

    private String packageFormData(MultiValueMap<String, String> formData,String charset){
        return formData.entrySet().stream().map(entry -> {
            String key = entry.getKey();
            String value = entry.getValue().stream().peek(v -> {
                try {
                    URLEncoder.encode(v, charset);
                } catch (UnsupportedEncodingException ignore) {
                }
            }).collect(Collectors.joining("&"));
            return String.format("%s=%s",key,value);
        }).collect(Collectors.joining("&"));
    }

    @Override
    public int getOrder() {
        return -9;
    }
}
